import argparse
import cli
import os
from parse_time import parse_time

def cmd_args():
	parser = argparse.ArgumentParser(description="Run WRF system.\n\nLongrun Weather Inc., NWP operation software.\nCopyright (C) 2018-2021 All Rights Reserved.", formatter_class=argparse.RawTextHelpFormatter)
	parser.add_argument('-c', '--codes', help='Root directory of all codes (e.g. WRF, WPS)')
	parser.add_argument(      '--wps-root', help='WPS root directory (e.g. WPS)')
	parser.add_argument(      '--wrf-root', help='WRF root directory (e.g. WRFV3 or WRF)')
	parser.add_argument(      '--wrfda-root', help='WRFDA root directory (e.g. WRFDA)')
	parser.add_argument(      '--wrfplus-root', help='WRFPLUS root directory (e.g. WRFPLUS)')    
	parser.add_argument(      '--upp-root', help='UPP root directory (e.g. UPP)')
	parser.add_argument(      '--met-root', help='MET root directory (e.g. MET)')
	parser.add_argument(      '--genbe-root', help='GENBE_2.0 root directory')
	parser.add_argument('-w', '--work-root', help='Work root directory')
	parser.add_argument('-g', '--geog-root', help='GEOG data root directory (e.g. WPS_GEOG)')
	parser.add_argument('-b', '--bkg-root', help='Background root directory')
	parser.add_argument('-r', '--ref-root', help='Reference root directory')
	parser.add_argument(      '--nature-run-root', help='Nature run root directory')
	parser.add_argument(      '--be-data', help='BE data file')
	parser.add_argument('-l', '--littler-root', help='LITTLE_R data root directory', default=None)
	parser.add_argument('-p', '--prepbufr-root', help='PrepBUFR data root directory')
	parser.add_argument('-j', '--config-json', help='Configuration JSON file.', required=True)
	parser.add_argument('-n', '--np', help='MPI process number to run WRF.', default=2, type=int)
	parser.add_argument(      '--ntasks-per-node', help='Override the default setting.', default=None, type=int)
	parser.add_argument(      '--node-list', help='Set node list for mpiexec.')
	parser.add_argument(      '--nodes', help='Set preferred nodes.', nargs='+')
	parser.add_argument(      '--excluded-nodes', help='Excluded bad nodes.', nargs='+')
	parser.add_argument(      '--exclusive', help='Set task to be exclusive in each node.', action='store_true')
	parser.add_argument(      '--slurm', help='Use SLURM job management system to run MPI jobs.', action='store_true')
	parser.add_argument(      '--pbs', help='Use PBS job management system variants (e.g. TORQUE) to run MPI jobs.', action='store_true')
	parser.add_argument('-t', '--start-time', help='Forecast start time (YYYYMMDDHH)', type=parse_time)
	parser.add_argument(      '--time', help='Process time (YYYYMMDDHH)', type=parse_time)
	parser.add_argument('-v', '--verbose', help='Print out work log', action='store_true')
	parser.add_argument('-f', '--force', help='Force to run', action='store_true')
	args = parser.parse_args()

	if not args.work_root:
		if os.getenv('WORK_ROOT'):
			args.work_root = os.getenv('WORK_ROOT')
		else:
			cli.error('Option --work-root or environment variable WORK_ROOT need to be set!')
	args.work_root = os.path.abspath(args.work_root)
	if not os.path.isdir(args.work_root):
		os.makedirs(args.work_root)
		cli.notice(f'Create work directory {args.work_root}.')

	if not args.wps_root:
		if os.getenv('WPS_ROOT'):
			args.wps_root = os.getenv('WPS_ROOT')
		elif args.codes:
			args.wps_root = args.codes + '/WPS'
		else:
			cli.error('Option --wps-root or environment variable WPS_ROOT need to be set!')
	args.wps_root = os.path.abspath(args.wps_root)
	if not os.path.isdir(args.wps_root):
		cli.error(f'Directory {args.wps_root} does not exist!')

	if not args.wrf_root:
		if os.getenv('WRF_ROOT'):
			args.wrf_root = os.getenv('WRF_ROOT')
		elif args.codes:
			args.wrf_root = args.codes + '/WRF'
		else:
			cli.error('Option --wrf-root or environment variable WRF_ROOT need to be set!')
	args.wrf_root = os.path.abspath(args.wrf_root)
	if not os.path.isdir(args.wrf_root):
		cli.error(f'Directory {args.wrf_root} does not exist!')

	if not args.wrfplus_root:
		if os.getenv('WRFPLUS_ROOT'):
			args.wrfplus_root = os.getenv('WRFPLUS_ROOT')
		elif args.codes:
			args.wrfplus_root = args.codes + '/WRFPLUS'
		else:
			cli.error('Option --wrfplus-root or environment variable WRFPLUS_ROOT need to be set!')
	args.wrfplus_root = os.path.abspath(args.wrfplus_root)
	if not os.path.isdir(args.wrfplus_root):
		cli.warning(f'Directory {args.wrfplus_root} does not exist!')

	if not args.wrfda_root:
		if os.getenv('WRFDA_ROOT'):
			args.wrfda_root = os.getenv('WRFDA_ROOT')
		elif args.codes:
			args.wrfda_root = args.codes + '/WRFDA'
		else:
			cli.error('Option --wrf-root or environment variable WRFDA_ROOT need to be set!')
	args.wrfda_root = os.path.abspath(args.wrfda_root)
	if not os.path.isdir(args.wrfda_root):
		cli.warning(f'Directory {args.wrfda_root} does not exist!')

	if not args.upp_root:
		if os.getenv('UPP_ROOT'):
			args.upp_root = os.getenv('UPP_ROOT')
		elif args.codes:
			args.upp_root = args.codes + '/UPP'
		else:
			cli.error('Option --upp-root or environment variable UPP_ROOT need to be set!')
	args.upp_root = os.path.abspath(args.upp_root)
	if not os.path.isdir(args.upp_root):
		cli.warning(f'Directory {args.upp_root} does not exist!')

	if not args.met_root:
		if os.getenv('MET_ROOT'):
			args.met_root = os.getenv('MET_ROOT')
		elif args.codes:
			args.met_root = args.codes + '/MET'
		else:
			cli.error('Option --met-root or environment variable MET_ROOT need to be set!')
	args.met_root = os.path.abspath(args.met_root)
	if not os.path.isdir(args.met_root):
		cli.warning(f'Directory {args.met_root} does not exist!')

	if not args.genbe_root:
		if os.getenv('GENBE_ROOT'):
			args.genbe_root = os.getenv('GENBE_ROOT')
		elif args.codes:
			args.genbe_root = args.codes + '/GENBE_2.0'
	if args.genbe_root:
		args.genbe_root = os.path.abspath(args.genbe_root)

	if not args.geog_root:
		if os.getenv('WPS_GEOG_ROOT'):
			args.geog_root = os.getenv('WPS_GEOG_ROOT')
		elif args.codes:
			args.geog_root = args.codes + '/WPS_GEOG'
		else:
			cli.error('Option --geog-root or environment variable WPS_GEOG_ROOT need to be set!')
	args.geog_root = os.path.abspath(args.geog_root)
	if not os.path.isdir(args.geog_root):
		cli.error(f'Directory {args.geog_root} does not exist!')

	if not args.bkg_root:
		if os.getenv('BKG_ROOT'):
			args.bkg_root = os.getenv('BKG_ROOT')
		else:
			cli.error('Option --bkg-root or environment variable BKG_ROOT need to be set!')
	args.bkg_root = os.path.abspath(args.bkg_root)
	if not os.path.isdir(args.bkg_root):
		cli.error(f'Directory {args.bkg_root} does not exist!')

	if args.nature_run_root:
		args.nature_run_root = os.path.abspath(args.nature_run_root)
		if not os.path.isdir(args.nature_run_root):
			cli.error(f'Directory {args.nature_run_root} does not exist!')

	if args.be_data:
		args.be_data = os.path.abspath(args.be_data)
		if not os.path.isfile(args.be_data):
			cli.error(f'File {args.be_data} does not exist!')

	return args
